import cv2
import threading, time
import numpy as np
from xgolib import XGO
from PIL import Image,ImageDraw,ImageFont
import xgoscreen.LCD_2inch as LCD_2inch
import os
from ultralytics import YOLO
import math

class SmartDogController():
    def __init__(self):
        self.dog = XGO("xgolite")
        self.display = LCD_2inch.LCD_2inch()
        self.display.Init() 
        self.display.clear()
        self.splash = Image.new("RGB",(320,240),"black")
        self.display.ShowImage(self.splash)
        self.cap = None
        self.camera_still = False
        self.save_path = "./saved_images"
        
        # 确保保存路径存在
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
        
        self.model = YOLO("best.pt")
        
        # 优化参数
        self.frame_skip = 2
        self.current_frame_count = 0
        
        # 位置调整参数
        self.target_center_x = 160  # 图像中心X坐标（理想夹取位置）
        self.target_center_y = 180  # 图像中心Y坐标（理想夹取位置）
        self.position_tolerance = 20  # 位置容差（像素）
        self.adjustment_active = False  # 是否正在调整位置
        self.last_detection_time = 0  # 上次检测到目标的时间
        
        # 调整参数
        self.move_speed = 10  # 移动速度
        self.turn_speed = 30  # 旋转速度
        self.arm_x = 100  # 机械臂X位置
        self.arm_z = -30  # 机械臂Z位置
        
        # 状态变量
        self.current_target = None  # 当前检测到的目标
        self.adjustment_history = []  # 调整历史
        
    def open_camera(self):
        print("📷 正在打开摄像头...")
        if self.cap is None:
            self.cap = cv2.VideoCapture(0)
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            self.cap.set(cv2.CAP_PROP_FPS, 30)
            print("✅ 摄像头初始化完成 (320x240@30fps, 低延迟模式)")
    
    def calculate_target_position(self, box):
        """计算目标在图像中的位置"""
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        center_x = (x1 + x2) // 2
        center_y = (y1 + y2) // 2
        width = x2 - x1
        height = y2 - y1
        area = width * height
        
        return {
            'center_x': center_x,
            'center_y': center_y,
            'width': width,
            'height': height,
            'area': area,
            'box': (x1, y1, x2, y2)
        }
    
    def analyze_position_error(self, target_pos):
        """分析位置误差并返回调整建议"""
        dx = target_pos['center_x'] - self.target_center_x
        dy = target_pos['center_y'] - self.target_center_y
        
        # 计算距离
        distance = math.sqrt(dx*dx + dy*dy)
        
        # 判断是否需要调整
        needs_adjustment = distance > self.position_tolerance
        
        # 计算调整方向
        if abs(dx) > 10:  # X方向调整阈值
            if dx > 0:
                x_adjustment = "right"  # 目标在右边，需要右移
            else:
                x_adjustment = "left"   # 目标在左边，需要左移
        else:
            x_adjustment = "center"
            
        if abs(dy) > 10:  # Y方向调整阈值
            if dy > 0:
                y_adjustment = "backward"  # 目标在后面，需要后退
            else:
                y_adjustment = "forward"   # 目标在前面，需要前进
        else:
            y_adjustment = "center"
        
        return {
            'needs_adjustment': needs_adjustment,
            'distance': distance,
            'dx': dx,
            'dy': dy,
            'x_adjustment': x_adjustment,
            'y_adjustment': y_adjustment,
            'target_size': target_pos['area']
        }
    
    def adjust_dog_position(self, adjustment_info):
        """根据调整信息移动机械狗"""
        if not adjustment_info['needs_adjustment']:
            print("🎯 位置已对准，可以夹取！")
            return True
            
        print(f"🔧 位置调整: X误差={adjustment_info['dx']:.1f}, Y误差={adjustment_info['dy']:.1f}")
        
        # 优先调整X方向（左右）
        if adjustment_info['x_adjustment'] == "right":
            print("➡️ 向右移动")
            self.dog.move_y(-self.move_speed)  # 右移
            time.sleep(0.5)
            self.dog.move_y(0)
        elif adjustment_info['x_adjustment'] == "left":
            print("⬅️ 向左移动")
            self.dog.move_y(self.move_speed)   # 左移
            time.sleep(0.5)
            self.dog.move_y(0)
        
        # 然后调整Y方向（前后）
        elif adjustment_info['y_adjustment'] == "forward":
            print("⬆️ 向前移动")
            self.dog.move_x(self.move_speed)   # 前进
            time.sleep(0.5)
            self.dog.move_x(0)
        elif adjustment_info['y_adjustment'] == "backward":
            print("⬇️ 向后移动")
            self.dog.move_x(-self.move_speed)  # 后退
            time.sleep(0.5)
            self.dog.move_x(0)
        
        return False
    
    def prepare_for_grasping(self):
        """准备夹取姿态"""
        print("🤖 调整夹取姿态...")
        
        # 蹲下姿态
        self.dog.translation('z', 85)  # 降低高度
        time.sleep(0.5)
        
        # 调整头部角度以便观察
        self.dog.motor(11, -30)  # 头部向下
        self.dog.motor(21, -30)  # 头部向下
        time.sleep(0.5)
        
        # 准备机械臂
        self.dog.arm(self.arm_x, self.arm_z)
        time.sleep(0.5)
        
        # 张开夹爪
        self.dog.claw(0)
        time.sleep(0.5)
        
        print("✅ 夹取姿态调整完成")
    
    def execute_grasp(self):
        """执行夹取动作"""
        print("🦾 开始夹取...")
        
        # 闭合夹爪
        self.dog.claw(255)
        time.sleep(1)
        
        # 抬起机械臂
        self.dog.arm(self.arm_x, -10)
        time.sleep(1)
        
        # 恢复站立姿态
        self.dog.translation('z', 105)
        time.sleep(1)
        
        # 复位
        self.dog.reset()
        time.sleep(2)
        
        print("✅ 夹取完成")
    
    def smart_camera_mode(self):
        """智能摄像头模式 - 包含位置调整逻辑"""
        if not self.cap.isOpened():
            print("无法打开摄像头")
            return
            
        print("🎯 智能检测模式启动")
        last_adjustment_time = 0
        
        while self.camera_still and self.cap is not None and self.cap.isOpened():
            # 获取最新帧
            ret = True
            frame = None
            for _ in range(3):
                ret, frame = self.cap.read()
                if not ret:
                    break
                    
            if not ret or frame is None:
                print("读取摄像头失败")
                time.sleep(0.01)
                continue
                
            self.current_frame_count += 1
            
            # 每隔frame_skip帧进行一次检测
            if self.current_frame_count % self.frame_skip == 0:
                # YOLO检测
                results = self.model.predict(frame, conf=0.3, verbose=False, imgsz=320)
                
                # 处理检测结果
                best_target = None
                best_confidence = 0
                
                for r in results:
                    if r.boxes is not None and len(r.boxes) > 0:
                        for box in r.boxes:
                            conf = float(box.conf[0])
                            cls_id = int(box.cls[0])
                            label = self.model.names[cls_id]
                            
                            # 只处理长方体目标
                            if label == "长方体" and conf > best_confidence:
                                best_confidence = conf
                                best_target = self.calculate_target_position(box)
                
                # 如果检测到目标
                if best_target is not None:
                    self.current_target = best_target
                    self.last_detection_time = time.time()
                    
                    # 分析位置误差
                    adjustment_info = self.analyze_position_error(best_target)
                    
                    # 在图像上绘制信息
                    x1, y1, x2, y2 = best_target['box']
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    
                    # 绘制目标中心点
                    cv2.circle(frame, (best_target['center_x'], best_target['center_y']), 5, (0, 0, 255), -1)
                    
                    # 绘制理想夹取位置
                    cv2.circle(frame, (self.target_center_x, self.target_center_y), 8, (255, 0, 0), 2)
                    
                    # 绘制连接线
                    cv2.line(frame, (best_target['center_x'], best_target['center_y']), 
                            (self.target_center_x, self.target_center_y), (255, 255, 0), 2)
                    
                    # 显示调整信息
                    info_text = f"DX:{adjustment_info['dx']:.0f} DY:{adjustment_info['dy']:.0f}"
                    cv2.putText(frame, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
                    
                    # 如果位置不对准且距离上次调整超过1秒，则进行调整
                    current_time = time.time()
                    if (adjustment_info['needs_adjustment'] and 
                        current_time - last_adjustment_time > 1.0):
                        
                        print(f"🎯 检测到长方体: 置信度={best_confidence:.2f}")
                        print(f"📍 位置: ({best_target['center_x']}, {best_target['center_y']})")
                        
                        # 执行位置调整
                        if self.adjust_dog_position(adjustment_info):
                            print("🎉 位置调整完成，准备夹取！")
                            # 这里可以添加自动夹取逻辑
                        
                        last_adjustment_time = current_time
                
                # 如果超过3秒没有检测到目标，显示提示
                elif time.time() - self.last_detection_time > 3.0:
                    cv2.putText(frame, "No target detected", (10, 60), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            
            # 显示处理后的图像
            try:
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb = cv2.flip(image_rgb, 1)
                pil_image = Image.fromarray(image_rgb)
                self.display.ShowImage(pil_image)
            except Exception as e:
                print(f"显示图像时出错: {e}")
                
            time.sleep(0.01)
    
    def auto_grasp_mode(self):
        """自动夹取模式"""
        print("🤖 启动自动夹取模式")
        
        # 准备夹取姿态
        self.prepare_for_grasping()
        
        # 启动智能检测
        self.smart_camera_mode()
    
    def manual_adjustment_mode(self):
        """手动调整模式 - 用于调试"""
        print("🔧 手动调整模式")
        print("按 'w' 前进, 's' 后退, 'a' 左移, 'd' 右移")
        print("按 'q' 左转, 'e' 右转, 'z' 蹲下, 'x' 站起")
        print("按 'g' 夹取, 'r' 复位, 'c' 退出")
        
        while True:
            try:
                key = input("请输入命令: ").lower()
                
                if key == 'w':
                    self.dog.move_x(10)
                    time.sleep(0.5)
                    self.dog.move_x(0)
                elif key == 's':
                    self.dog.move_x(-10)
                    time.sleep(0.5)
                    self.dog.move_x(0)
                elif key == 'a':
                    self.dog.move_y(10)
                    time.sleep(0.5)
                    self.dog.move_y(0)
                elif key == 'd':
                    self.dog.move_y(-10)
                    time.sleep(0.5)
                    self.dog.move_y(0)
                elif key == 'q':
                    self.dog.turn(30)
                    time.sleep(0.5)
                    self.dog.turn(0)
                elif key == 'e':
                    self.dog.turn(-30)
                    time.sleep(0.5)
                    self.dog.turn(0)
                elif key == 'z':
                    self.dog.translation('z', 85)
                elif key == 'x':
                    self.dog.translation('z', 105)
                elif key == 'g':
                    self.execute_grasp()
                elif key == 'r':
                    self.dog.reset()
                elif key == 'c':
                    break
                    
            except KeyboardInterrupt:
                break
        
        print("手动调整模式结束")
    
    def xgoCamera(self, switch):
        print(f"🎮 摄像头开关: {'开启' if switch else '关闭'}")
        if switch:
            self.open_camera()
            if self.cap and self.cap.isOpened():
                self.camera_still = True
                self.current_frame_count = 0
                
                # 启动智能检测线程
                t = threading.Thread(target=self.smart_camera_mode, daemon=True)  
                t.start() 
                
                # 调整摄像头角度
                print("🔧 调整机械狗摄像头角度...")
                self.dog.motor(11, -30)
                self.dog.motor(21, -30)
                print("✅ 摄像头角度调整完成")
            else:
                print("❌ 摄像头打开失败")
        else:
            print("🛑 正在关闭摄像头...")
            self.camera_still = False
            time.sleep(0.1)
            if self.cap:
                self.cap.release()
                self.cap = None
            splash = Image.new("RGB", (320, 240), "black")
            self.display.ShowImage(splash)
            print("✅ 摄像头已关闭")

    def close(self):
        """清理资源"""
        print("🧹 正在清理系统资源...")
        self.camera_still = False
        time.sleep(0.2)
        if self.cap:
            self.cap.release()
            print("✅ 摄像头资源已释放")
        self.display.clear()
        print("✅ 显示屏已清理")
        print("💤 系统关闭完成")

# 使用示例
if __name__ == "__main__":
    print("🤖 智能机械狗位置调整系统")
    print("=" * 60)
    print("📌 功能特性:")
    print("   • 实时目标检测")
    print("   • 智能位置调整")
    print("   • 自动夹取准备")
    print("   • 手动调试模式")
    print("📌 按 Ctrl+C 退出程序")
    print("=" * 60)
    
    x = SmartDogController()
    try:
        # 选择模式
        print("请选择模式:")
        print("1. 智能检测模式")
        print("2. 自动夹取模式")
        print("3. 手动调整模式")
        
        choice = input("请输入选择 (1-3): ").strip()
        
        if choice == "1":
            x.xgoCamera(True)
            while True:
                time.sleep(0.1)
        elif choice == "2":
            x.xgoCamera(True)
            time.sleep(2)  # 等待摄像头稳定
            x.auto_grasp_mode()
        elif choice == "3":
            x.manual_adjustment_mode()
        else:
            print("无效选择，启动智能检测模式")
            x.xgoCamera(True)
            while True:
                time.sleep(0.1)
            
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断程序")
        print("🛑 正在安全关闭系统...")
        x.xgoCamera(False)
    finally:
        x.close()
        print("\n🎯 程序运行结束")
        print("=" * 60) 